Skip to content

Conversation

pctablet505
Copy link
Collaborator

@pctablet505 pctablet505 commented Sep 17, 2025

This pull request adds support for exporting Keras models to the LiteRT (TFLite) format, along with improvements to input signature handling and export utility documentation. The changes ensure that LiteRT export is only available when TensorFlow is installed, update the export API and documentation, and enhance input signature inference for various model types.

LiteRT Export Support:

  • Added conditional import of LitertExporter and export_litert in keras/src/export/__init__.py, making LiteRT export available only if TensorFlow is installed.
  • Updated the Model.export method to support the "litert" format, including new options for LiteRT export and user-facing documentation and example. Raises an informative error if TensorFlow is not installed. [1] [2] [3] [4]
  • Registered litert as a lazy module in keras/src/utils/module_utils.py for dynamic import support.

Input Signature and Export Utilities:

  • Improved documentation and logic in get_input_signature to clarify behavior for different model types and ensure correct input signature construction for export. [1] [2]
  • Enhanced _infer_input_signature_from_model to handle flexible batch dimensions and ensure compatibility with downstream exporters, always returning a flat list of input specs.

pctablet505 and others added 30 commits August 28, 2025 11:04
Introduces a custom LiteRTExporter for exporting models to TFLite format, bypassing the standard TFLiteConverter. Updates the export API and documentation to support the new 'lite_rt' format, and adds relevant options for custom ops, select TF ops, and optimizations.
Replaces the custom MLIR-based TFLite conversion logic in LiteRTExporter with direct use of the standard TFLiteConverter. Also improves input signature handling for tf.function tracing and updates imports accordingly.
Moved imports of get_input_signature and make_tf_tensor_spec inside functions in saved_model.py to prevent circular imports. Updated EXPORT_FORMATS in export_utils.py to use string references instead of direct imports.
Adds max_sequence_length parameter to input signature generation for sequence models, bounding sequence length for transformer-like architectures. Improves LiteRTExporter with heuristics for complex models, fallback conversion via SavedModel for large models, and exposes max_sequence_length in Model export options. Updates documentation accordingly.
Adds logic to dynamically reduce max_sequence_length for large vocabulary models in export_utils.py to prevent tensor size overflow. In lite_rt_exporter.py, introduces checks and workarounds for models with _DictWrapper issues, and applies memory optimizations for large models during TFLite conversion. These changes improve export reliability and prevent memory errors for models such as Gemma, Llama, and similar architectures.
Removed custom trackable object logic from LiteRTExporter and now save the model directly, simplifying the export process. Also streamlined vocabulary size checks in export_utils to prevent tensor size overflow, removing verbose warnings and redundant comments.
Refactors the TFLite conversion logic in lite_rt_exporter.py to attempt direct conversion first and only fall back to SavedModel if necessary, improving robustness and clarity. Adds a new lite_rt_exporter_simple.py file with a streamlined LiteRTExporter class for direct TFLite export, bypassing complex MLIR conversion paths.
Refactors export_utils and lite_rt_exporter to better detect large vocabulary and Keras-Hub models, applying safer sequence length limits and more robust TFLite conversion paths. Adds heuristics for model type detection, ensures memory safety, and improves handling of TensorFlow introspection issues during export.
Working well with keras
Eliminates the logic for bounding sequence length in model export utilities and related code paths. The max_sequence_length parameter and associated shape bounding for large vocabulary models are removed from export_utils.py and lite_rt_exporter.py. Updates model export documentation accordingly. Adds a comprehensive test script for Keras Hub LiteRT export, verifying numerical accuracy between original and exported models.
Renames the exporter module and class from 'lite_rt_exporter' and 'LiteRTExporter' to 'litert_exporter' and 'LitertExporter', respectively. Updates all references, registry keys, and documentation to use 'litert' instead of 'lite_rt'. Adds the new 'export_litert' function and corresponding tests for Litert export functionality.
Refactors and improves the LiteRT exporter for better input signature inference, error handling, and verbose output. Updates and expands the test suite to cover multi-input, multi-output, and error scenarios, and ensures the Model.export() API is used consistently. Also improves docstrings and formatting for clarity and maintainability.
Replaces direct imports and availability checks for LiteRT with the LazyModule-based `litert` utility from `module_utils`. Updates tests and exporter logic to require '.tflite' filepaths and consistently use the new import pattern, improving modularity and error handling.
@pctablet505
Copy link
Collaborator Author

@gemini-code-assist

Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request introduces support for exporting Keras models to the LiteRT (TFLite) format, which is a valuable addition for mobile and edge deployment. The changes are well-structured, including a new LitertExporter, updates to the Model.export API and documentation, and related refactoring of export utilities. The implementation appears robust, with thoughtful handling of different model types and conversion strategies, and is accompanied by a comprehensive test suite. My feedback focuses on a few areas to improve code clarity, maintainability, and adherence to best practices.

Copy link
Contributor

Hello @pctablet505! Thanks for reaching out. I'm here to assist with your pull request, "Model Export to liteRT".

I noticed from the Codecov report that the new changes have low patch coverage (10.20408%). This is primarily due to missing coverage in keras/src/export/litert_exporter.py and keras/src/export/export_utils.py.

If you have any questions about improving the test coverage, clarifying specific parts of the code, or anything else related to this PR, please let me know how I can help!

pctablet505 and others added 10 commits October 1, 2025 12:13
Removed the return value and conditional messaging from export_litert. The function now always prints a simple message upon export and no longer returns the compilation result or filepath.
Simplifies and standardizes input signature generation for Functional, Sequential, and subclassed models to always use a flexible batch dimension (None) for export. Refactors Sequential input shape inference to remove hardcoded fallbacks and ensure dynamic batching support. Cleans up and clarifies error handling and structure of input signature utilities.
Renamed litert_exporter.py and its test to litert.py and litert_test.py, respectively. Updated all relevant imports to use the new module name. Refactored LitertExporter to simplify model building logic and removed redundant input shape inference for Sequential models. Also cleaned up export_utils.py by removing the export format registry and related functions.
Moved imports of get_input_signature and make_tf_tensor_spec to the module level in saved_model.py for consistency and efficiency. In litert.py, imports remain local to preserve lazy loading where needed.
@pctablet505
Copy link
Collaborator Author

@gemini-code-assist

Copy link
Contributor

Hello @pctablet505! Thanks for reaching out again. I'm still here to help with your pull request, "Model Export to liteRT".

My previous comment highlighted the low patch coverage (12.73585%) in keras/src/export/litert.py and keras/src/export/export_utils.py as reported by Codecov.

Would you like me to provide some specific suggestions on how to improve the test coverage for these files, or do you have any other questions about the PR that I can assist with?

LiteRT export and related tests now import TensorFlow conditionally, allowing Keras to function without TensorFlow installed. If TensorFlow is unavailable, LiteRT export is disabled and tests are skipped, with informative messages and error handling added to guide users.
@pctablet505 pctablet505 requested a review from abheesht17 October 6, 2025 05:16
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants